import numpy as np
import matplotlib.pyplot as plt
import random
import time

# ================== 参数配置区 ==================
# 用户可自由修改以下参数
CITY_NUM = 20          # 城市数量（问题规模）
POP_SIZE = 100         # 种群规模（个体数量）
MAX_GENERATIONS = 1000  # 最大迭代次数（进化代数）
MUTATION_RATE = 0.2    # 变异概率（控制种群多样性）
CROSSOVER_RATE = 0.8   # 交叉概率（控制基因重组频率）
ELITE_SIZE = 10        # 精英个体数量（直接保留到下一代，保证最优解不丢失）


# ================== 城市生成与距离计算 ==================
def generate_cities(num_cities):
    """生成随机城市坐标
    遗传算法解决TSP问题的第一步：定义问题空间
    每个城市用(x, y)坐标表示，坐标范围[0,100]
    """
    np.random.seed(100)  # 固定随机种子，确保结果可复现
    return np.random.rand(num_cities, 2) * 100


def calculate_distance(city1, city2):
    """计算两个城市之间的欧几里得距离：直线距离
    这是TSP问题的核心度量标准
    """
    return np.sqrt((city1[0] - city2[0]) ** 2 + (city1[1] - city2[1]) ** 2)


def calculate_total_distance(path, cities):
    """计算路径总距离
    适应度函数：评估个体（路径）的优劣
    距离越短表示适应度越高
    """
    total_distance = 0
    num_cities = len(path)

    for i in range(num_cities):
        from_city = path[i]
        to_city = path[(i + 1) % num_cities]  # 循环回到起点（TSP是闭环问题）
        total_distance += calculate_distance(cities[from_city], cities[to_city])

    return total_distance


# ================== 遗传算法操作 ==================
def initialize_population(pop_size, num_cities):
    """初始化种群：生成随机路径
    遗传算法起点：创建初始解集
    每个个体代表一条访问所有城市的随机路径
    """
    population = []
    for _ in range(pop_size):
        # 生成0到num_cities-1的随机排列
        path = list(range(num_cities))
        random.shuffle(path)  # 随机打乱顺序
        population.append(path)
    return population


def rank_population(population, cities):
    """评估种群并排序（按路径长度升序）
    选择操作的前置步骤：评估种群中每个个体的适应度
    返回排序后的列表：(个体索引, 路径长度)
    """
    fitness_results = []
    for i, path in enumerate(population):
        distance = calculate_total_distance(path, cities)
        fitness_results.append((i, distance))

    # 按路径长度排序（距离越短越好）
    return sorted(fitness_results, key=lambda x: x[1])


# ================== 选择操作 ==================
def roulette_wheel_selection(population, ranked_population, elite_size):
    """轮盘赌选择
    基于适应度比例的选择方法
    适应度高的个体有更高概率被选中
    """
    selection_results = []

    # 1. 首先保留精英个体（精英策略：直接保留最优个体）
    for i in range(elite_size):
        selection_results.append(population[ranked_population[i][0]])

    # 2. 计算适应度值（距离的倒数：距离越短，适应度越高）
    fitness_values = [1 / ranked_population[i][1] for i in range(len(ranked_population))]
    total_fitness = sum(fitness_values)

    # 3. 计算选择概率（概率与适应度成正比）
    selection_probs = [f / total_fitness for f in fitness_values]

    # 4. 使用轮盘赌选择剩余个体
    for _ in range(len(population) - elite_size):
        pick = random.random()  # 生成[0,1)随机数
        current = 0
        for i in range(len(ranked_population)):
            current += selection_probs[i]
            if current > pick:  # 当累积概率超过随机数时选择该个体
                selection_results.append(population[ranked_population[i][0]])
                break

    return selection_results


def tournament_selection(population, ranked_population, elite_size, tournament_size=5):
    """锦标赛选择
    另一种选择策略：随机选取若干个体竞争，最优者胜出
    参数tournament_size控制选择压力（值越大，选择压力越大）
    """
    selection_results = []

    # 1. 首先保留精英个体
    for i in range(elite_size):
        selection_results.append(population[ranked_population[i][0]])

    # 2. 锦标赛选择剩余个体
    for _ in range(len(population) - elite_size):
        # 随机选择tournament_size个个体进行"比赛"
        tournament = random.sample(ranked_population, tournament_size)
        # 选择距离最短的个体（适应度最高）
        winner = min(tournament, key=lambda x: x[1])
        selection_results.append(population[winner[0]])

    return selection_results


# ================== 交叉操作 ==================
def ordered_crossover(parent1, parent2):
    """有序交叉(OX)
    常用TSP交叉方法：保留父代1的片段，按父代2顺序填充剩余城市
    特点：保持城市访问顺序，避免非法路径
    """
    size = len(parent1)

    # 选择两个随机交叉点（确定保留的片段）
    start, end = sorted(random.sample(range(size), 2))

    # 初始化子代
    child = [None] * size

    # 从parent1复制片段（保留父代1的优良基因）
    child[start:end] = parent1[start:end]

    # 从parent2填充剩余部分（按父代2的顺序）
    current_pos = end % size
    for gene in parent2:
        if gene not in child:  # 只添加未包含的城市
            child[current_pos] = gene
            current_pos = (current_pos + 1) % size  # 循环填充

    return child


def partially_mapped_crossover(parent1, parent2):
    """部分映射交叉(PMX)
    另一种常用TSP交叉方法：通过映射关系解决冲突
    特点：保持城市唯一性，适合路径表示
    """
    size = len(parent1)

    # 选择两个随机交叉点
    start, end = sorted(random.sample(range(size), 2))

    # 初始化子代
    child1 = [None] * size
    child2 = [None] * size

    # 复制交叉段（直接交换父代片段）
    child1[start:end] = parent1[start:end]
    child2[start:end] = parent2[start:end]

    # 创建映射关系（解决城市重复问题）
    mapping1 = {}  # parent2到parent1的映射
    mapping2 = {}  # parent1到parent2的映射

    for i in range(start, end):
        # 建立parent2[i] -> parent1[i]的映射
        if parent2[i] not in child1[start:end]:
            mapping1[parent2[i]] = parent1[i]
        # 建立parent1[i] -> parent2[i]的映射
        if parent1[i] not in child2[start:end]:
            mapping2[parent1[i]] = parent2[i]

    # 填充子代1的剩余位置
    for i in range(size):
        if child1[i] is None:
            gene = parent2[i]
            # 处理映射链（直到找到非冲突基因）
            while gene in mapping1:
                gene = mapping1[gene]
            child1[i] = gene

    # 填充子代2的剩余位置
    for i in range(size):
        if child2[i] is None:
            gene = parent1[i]
            while gene in mapping2:
                gene = mapping2[gene]
            child2[i] = gene

    return child1, child2


# ================== 变异操作 ==================
def swap_mutation(path):
    """交换变异：随机交换两个城市的位置
    增加种群多样性，避免早熟收敛
    """
    if random.random() < MUTATION_RATE:  # 以MUTATION_RATE概率执行变异
        # 随机选择两个不同位置
        idx1, idx2 = random.sample(range(len(path)), 2)
        # 交换位置
        path[idx1], path[idx2] = path[idx2], path[idx1]
    return path


def inversion_mutation(path):
    """逆序变异：随机选择一段路径并反转
    局部优化技术，可能改善路径质量
    """
    if random.random() < MUTATION_RATE:
        # 随机选择片段起止点
        start, end = sorted(random.sample(range(len(path)), 2))
        # 反转片段
        path[start:end] = reversed(path[start:end])
    return path


def scramble_mutation(path):
    """打乱变异：随机选择一段路径并打乱
    更强的多样性引入方式
    """
    if random.random() < MUTATION_RATE:
        start, end = sorted(random.sample(range(len(path)), 2))
        segment = path[start:end]
        random.shuffle(segment)  # 随机打乱片段
        path[start:end] = segment
    return path


# ================== 遗传算法主函数 ==================
def genetic_algorithm(cities, selection_method=tournament_selection,
                      crossover_method=ordered_crossover, mutation_method=swap_mutation):
    """遗传算法主函数
    实现完整的进化流程：
    初始化 → 评估 → 选择 → 交叉 → 变异 → 替换 → 重复
    """
    # 记录每代最佳总距离（用于分析收敛过程）
    progress = []

    # 1. 初始化种群
    population = initialize_population(POP_SIZE, len(cities))

    # 2. 评估初始种群
    ranked_population = rank_population(population, cities)
    progress.append(ranked_population[0][1])  # 记录最佳距离

    # 3. 进化过程（核心循环）
    for generation in range(MAX_GENERATIONS):
        # 3.1 选择：从当前种群中选择优秀个体作为父代
        selected = selection_method(population, ranked_population, ELITE_SIZE)

        # 3.2 交叉：重组父代基因创建子代
        children = []
        # 首先保留精英（直接进入下一代）
        children.extend(selected[:ELITE_SIZE])

        # 交叉生成后代（成对处理父代）
        for i in range(ELITE_SIZE, len(selected), 2):
            parent1 = selected[i]
            parent2 = selected[i + 1]

            # 以CROSSOVER_RATE概率执行交叉
            if random.random() < CROSSOVER_RATE:
                # 根据交叉方法生成子代
                if crossover_method == partially_mapped_crossover:
                    # PMX交叉返回两个子代
                    child1, child2 = crossover_method(parent1, parent2)
                    children.append(child1)
                    children.append(child2)
                else:
                    # OX交叉每次生成一个子代，顺序交叉
                    child = crossover_method(parent1, parent2)
                    children.append(child)
                    # 对称生成另一个子代
                    children.append(crossover_method(parent2, parent1))
            else:
                # 不交叉，直接复制父代（保持种群多样性）
                children.append(parent1)
                children.append(parent2)

        # 3.3 变异：引入随机变化随机交换两个城市的循序
        for i in range(len(children)):
            children[i] = mutation_method(children[i])

        # 3.4 更新种群：用子代替换当前种群
        population = children

        # 3.5 评估新种群
        ranked_population = rank_population(population, cities)
        progress.append(ranked_population[0][1])  # 记录当前代最优解

        # 3.6 打印进度（每50代输出一次）
        if generation % 50 == 0:
            print(f"Generation {generation}: Best Distance = {ranked_population[0][1]:.2f}")

    # 4. 返回最终结果
    best_index = ranked_population[0][0]
    best_path = population[best_index]
    best_distance = ranked_population[0][1]

    print(f"\nFinal Result: Best Distance = {best_distance:.2f}")
    return best_path, best_distance, progress


# ================== 可视化函数 ==================
def plot_cities(cities):
    """绘制城市分布图
    可视化问题空间
    """
    plt.figure(figsize=(10, 6))
    x = [city[0] for city in cities]
    y = [city[1] for city in cities]
    plt.scatter(x, y, c='red', marker='o')

    # 标注城市编号
    for i, city in enumerate(cities):
        plt.annotate(str(i), (city[0], city[1]), fontsize=12)

    plt.title(f'City Distribution ({len(cities)} Cities)')
    plt.xlabel('X Coordinate')
    plt.ylabel('Y Coordinate')
    plt.grid(True)
    plt.show()


def plot_path(path, cities, title):
    """绘制路径图
    可视化解决方案
    """
    plt.figure(figsize=(10, 6))
    x = [city[0] for city in cities]
    y = [city[1] for city in cities]
    plt.scatter(x, y, c='red', marker='o')

    # 绘制路径连线
    for i in range(len(path)):
        start_city = cities[path[i]]
        end_city = cities[path[(i + 1) % len(path)]]
        plt.plot([start_city[0], end_city[0]], [start_city[1], end_city[1]], 'b-')

    # 标记起点
    start_city = cities[path[0]]
    plt.annotate('Start', (start_city[0], start_city[1]), fontsize=12, color='green')

    plt.title(f'{title} (Distance: {calculate_total_distance(path, cities):.2f})')
    plt.xlabel('X Coordinate')
    plt.ylabel('Y Coordinate')
    plt.grid(True)
    plt.show()


def plot_progress(progress, title):
    """绘制进化过程
    可视化算法收敛情况
    """
    plt.figure(figsize=(10, 6))
    plt.plot(progress)
    plt.title(title)
    plt.xlabel('Generation')
    plt.ylabel('Best Distance')
    plt.grid(True)
    plt.show()


# ================== 主程序 ==================
if __name__ == "__main__":
    # 生成城市（问题实例）
    cities = generate_cities(CITY_NUM)
    plot_cities(cities)  # 可视化城市分布

    # 打印算法参数
    print(f"=== 遗传算法参数 ===")
    print(f"城市数量: {CITY_NUM}")
    print(f"种群大小: {POP_SIZE}")
    print(f"最大代数: {MAX_GENERATIONS}")
    print(f"变异概率: {MUTATION_RATE}")
    print(f"交叉概率: {CROSSOVER_RATE}")
    print(f"精英数量: {ELITE_SIZE}")

    # 运行遗传算法
    start_time = time.time()
    best_path, best_distance, progress = genetic_algorithm(cities)
    end_time = time.time()

    # 输出结果
    print(f"\n运行时间: {end_time - start_time:.2f}秒")
    print(f"最优路径: {best_path}")

    # 可视化结果
    plot_path(best_path, cities, "Optimal Path")  # 最优路径可视化
    plot_progress(progress, "Genetic Algorithm Progress")  # 收敛过程可视化